Skip to content

[python] Add direct LeRobot capture writer - #9665

Open
YannByron wants to merge 2 commits into
apache:masterfrom
YannByron:codex/lerobot-direct-writer
Open

[python] Add direct LeRobot capture writer#9665
YannByron wants to merge 2 commits into
apache:masterfrom
YannByron:codex/lerobot-direct-writer

Conversation

@YannByron

@YannByron YannByron commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Summary

Add a LeRobot recording-loop compatible writer that sends captured frames directly to a Paimon multimodal table without materializing an intermediate LeRobot dataset.

This PR builds on the generic PyPaimon Snapshot-properties commit support merged in #9692 and uses it to persist LeRobot resume state.

Changes

  • Add PaimonLeRobotWriter with the recording lifecycle used by LeRobot: frame buffering, episode save/discard, pending-frame checks, and finalize.
  • Match native frame requirements: all user features plus task, generated indices excluded, exact numeric dtype/shape, and CHW/HWC image support.
  • Default to one Paimon commit at finalize(); optional positive episodes_per_commit and flush() provide explicit intermediate commit boundaries.
  • Append to compatible existing tables after exact schema validation.
  • Persist next frame, episode, and task indices in Snapshot properties, avoiding table scans on normal resume; legacy tables without state are scanned once.
  • Keep clear_episode_buffer() limited to the current unsaved episode; already accepted episodes remain in the pending batch.
  • Encode image features as PNG BLOBs; MP4/video features remain unsupported in this first version.

Testing

  • python -m pytest -q pypaimon/tests/multimodal_lerobot_test.py pypaimon/tests/multimodal_lerobot_writer_test.py pypaimon/tests/table_commit_test.py pypaimon/tests/table/simple_table_test.py (92 passed)
  • Ruff checks for all changed Python files
  • git diff --check

Notes

This draft intentionally starts with PNG image frames, synchronous encoding, and a single active writer per table. MP4 encoding, concurrent index allocation, and streaming-checkpoint semantics are deferred.

@YannByron
YannByron force-pushed the codex/lerobot-direct-writer branch from 14eea28 to c7b0a9d Compare September 8, 2026 12:59
@YannByron
YannByron marked this pull request as ready for review September 8, 2026 14:25
Add a LeRobot-compatible writer that batches completed episodes into multimodal Paimon commits, resumes compatible tables, and stores image frames as PNG BLOBs.

Co-Authored-By: Codex <noreply@anthropic.com>

AI-Model: gpt-5
AI-Contributed/Feature: 424/424
AI-Contributed/UT: 366/366
Allow table commits to attach snapshot properties and use them to resume LeRobot frame, episode, and task indices without scanning table data.

Co-Authored-By: Codex <noreply@anthropic.com>

AI-Model: gpt-5
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>

AI-Contributed/Feature: 71/71
AI-Contributed/UT: 40/40
@YannByron
YannByron force-pushed the codex/lerobot-direct-writer branch from b4aedae to 0c3ebdd Compare September 8, 2026 14:28

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found two reproducible input-compatibility issues in the direct capture writer; details are inline.

Comment on lines +273 to +275
return _encode_media_frame(
value,
channel_first=actual_shape == expected_shape,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Preserve the layout of native LeRobot HWC camera features

Matching the declared shape does not imply CHW. LeRobot 0.4.4's hw_to_dataset_features preserves the camera's HWC shape and sets names=["height", "width", "channels"]. With a standard image feature declaring shape=(480, 640, 3) and a matching np.uint8 frame, this sets channel_first=True and transposes the image to (640, 3, 480). I reproduced the first add_frame() failing with Unsupported LeRobot media frame shape or dtype: (640, 3, 480), uint8, so the advertised native recording-loop integration cannot record standard camera frames. Please determine the layout from the feature's dimension names and actual channel axis, and add a regression test using native HWC feature metadata.

raise ValueError(
"LeRobot feature %s expected shape %s, got %s."
% (name, expected_shape, value.shape))
return _normalize_value(value, feature, name)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Normalize validated NumPy arrays before passing them to the Arrow helpers

For non-scalar features, _normalize_value() returns the original ndarray, whereas the existing import path supplies Python values from to_pylist(). This makes valid native LeRobot inputs fail in the subsequent _safe_array([value], ...): a float32 array with shape (2, 2) raises Can only convert 1-dimensional array values, and a bool array with shape (2,) is rejected because its np.bool_ elements are not Python bool instances. Both satisfy LeRobot's native dtype/shape validation, and I reproduced both failures through PaimonLeRobotWriter.add_frame(). Please convert validated arrays to Python scalars/nested lists before buffering and Arrow validation, and cover both cases. Applying .tolist() after validation allowed both examples to complete a write/read round trip.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants